# Copyright (c) Alibaba, Inc. and its affiliates.
import datetime as dt
import os
from abc import ABC, abstractmethod
from typing import List, Union

from swift.utils import get_logger, parse_args, seed_everything
from .argument import BaseArguments
from .utils import ProcessorMixin

logger = get_logger()


class SwiftPipeline(ABC, ProcessorMixin):
    # 指定默认使用的参数类（这里是 BaseArguments，可被子类覆盖）
    args_class = BaseArguments

    # 初始化函数，接收命令行参数或配置类实例
    def __init__(self, args: Union[List[str], args_class, None] = None):
        # 将参数解析为 args_class 实例
        self.args = self._parse_args(args)
        args = self.args

        # 如果包含种子 seed，则设置全局随机种子（可支持分布式 rank 偏移）
        if hasattr(args, 'seed'):
            seed = args.seed + max(getattr(args, 'rank', -1), 0)
            seed_everything(seed)  # 通常设置 torch、numpy、random 的种子

        # 打印解析后的参数信息
        logger.info(f'args: {args}')

        # 处理在 Data Science Workshop 环境下运行 Gradio 的兼容性
        self._compat_dsw_gradio(args)

    # ---------------- 参数解析逻辑 ---------------- #
    def _parse_args(self, args: Union[List[str], args_class, None] = None) -> args_class:
        # 如果传入的是 BaseArguments 实例，直接返回
        if isinstance(args, self.args_class):
            return args

        assert self.args_class is not None  # 确保参数类已定义

        # 使用 parse_args 将命令行参数转换为 args_class 实例
        args, remaining_argv = parse_args(self.args_class, args)

        # 如果存在未识别的参数，根据配置判断是否报错
        if len(remaining_argv) > 0:
            if getattr(args, 'ignore_args_error', False):
                logger.warning(f'remaining_argv: {remaining_argv}')  # 仅记录警告
            else:
                raise ValueError(f'remaining_argv: {remaining_argv}')  # 抛出错误
        return args

    # ---------------- Gradio + DSW 环境兼容处理 ---------------- #
    @staticmethod
    def _compat_dsw_gradio(args) -> None:
        # 导入 Web UI 参数类
        from swift.llm import WebUIArguments, AppArguments

        # 检查是否运行在 Data Science Workshop（DSW）环境中
        if (isinstance(args, (WebUIArguments, AppArguments)) and
                'JUPYTER_NAME' in os.environ and
                'dsw-' in os.environ['JUPYTER_NAME'] and
                'GRADIO_ROOT_PATH' not in os.environ):
            # 自动为 Gradio 设置代理路径（用于前端访问）
            os.environ['GRADIO_ROOT_PATH'] = f"/{os.environ['JUPYTER_NAME']}/proxy/{args.server_port}"

    # ---------------- 运行主流程的入口 ---------------- #
    def main(self):
        logger.info(f'Start time of running main: {dt.datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")}')
        result = self.run()  # 调用 run 方法（需子类实现）
        logger.info(f'End time of running main: {dt.datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f")}')
        return result

    # ---------------- 抽象方法：主运行逻辑 ---------------- #
    @abstractmethod
    def run(self):
        # 抽象方法，子类必须实现
        pass
